
import base64
from zhipuai import ZhipuAI
import os
image_dir = "/Users/muzhiyi/Desktop/perscription_senior/perscription-basic"
image_files = [f for f in os.listdir(image_dir) if f.endswith(('.png', '.jpg',  '.jpeg'))]
import csv

# CSV 文件路径
csv_file_path = '/Users/muzhiyi/Desktop/perscription_senior/output_basic/base_glm-4v-flash_output.csv'

if not os.path.exists(csv_file_path):
    with open(csv_file_path, 'w', newline='', encoding='utf-8') as csvfile:
        fieldnames = ['filename', 'description', 'analysis']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
else:
    pass

def call_api_picture_analysis(img_path):
    with open(img_path, 'rb') as img_file:
        img_base = base64.b64encode(img_file.read()).decode('utf-8')

    client = ZhipuAI(api_key="cb90d533fd12a70a36e6a21e7371e092.qtVWybqNm4GCqg07")  # 填写您自己的APIKey
    response = client.chat.completions.create(
        model="glm-4v-flash",  # 填写需要调用的模型名称
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": img_base
                        }
                    },
                    {
                        "type": "text",
                        "text":"识别图中处方内容， 要求给出{{姓名：}，{性别：},{年龄：},{科室：},{时间：},{诊断：},{药物1名称： ,用法：},{药物2名称: ， 用法： }...,{医师：},{药师：},{审核： },{调配： },{核对： },{发药： },{金额： }}等信息"
                    }
                ]
            }
        ]
    )
    output1 = response.choices[0].message.content
    # print(response.choices[0].message)
    print("模型第一轮输出 \n")
    print(output1)

    response2 = client.chat.completions.create(
        model="glm-4v-flash",  # 填写需要调用的模型名称
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": img_base
                        }
                    },
                    {
                        "type": "text",
                        "text":"识别图中处方内容， 要求给出{{姓名：}，{性别：},{年龄：},{科室：},{时间：},{诊断：},{药物1名称： ,用法：},{药物2名称: ， 用法： }.（其他药物写法以此类推）..,{医师：},{药师：},{审核： },{调配： },{核对： },{发药： },{金额： }}等信息"
                    }
                ]
            },
            {
                "content": [
                    {
                        "text": output1,
                        "type": "text"
                    }
                ],
                "role": "assistant"
            },
            {
                "role": "user",
                "content":[
                    {
                        "type":"text",
                        "text": "关于给出的处方内容，结合现有的医学知识，识别药物搭配是否合理，是否存在剂量过多的情况。 如果是安全的，则回答safe， 如果不合理， 则回答unsafe，并给出相应原因。回答格式为：先结论：Safe/Unsafe。 后原因：Reason:....."
                    }
                ]
            }
        ]
    )
    output2 = response2.choices[0].message.content
    print("模型第二轮输出 \n")
    print(output2)


    return output1,output2

# 后续的写入操作可以使用 'a'（追加）模式或 'w'（覆盖）模式
with open(csv_file_path, 'a', newline='', encoding='utf-8') as csvfile:
    fieldnames = ['filename', 'description', 'analysis']
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

    for image_file in image_files:
        image_path = os.path.join(image_dir, image_file)
        output1, output2 = call_api_picture_analysis(image_path)
        writer.writerow({'filename': image_file, 'description': output1, 'analysis': output2})
